import numpy as np
import PIL.Image as Image
import patchreconst as pr
import matplotlib.pyplot as plt

PATCH_SIZE = 32
COMPONENT_COUNT = 100
PATCH_COUNT = 400000
EPOCH_COUNT = 16
CODE_COUNT = 800
SPARSITY_PARAMETER = 4.0
NON_NEGATIVE = True

scBasisFile = 'basis/basisnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'
scBasis = np.load(scBasisFile)

icaFilterFile = 'icabasis/basisluisicacomps' + str(COMPONENT_COUNT) + 'codes' + str(CODE_COUNT) + 'patches' + str(PATCH_COUNT) + '.npy'
icaFilters = np.load(icaFilterFile)

patches = np.load('patches_1000_mod.npy')
patchesOrig = np.load('patches_1000.npy')[:, :, :, 0]
rmBounds = np.load('patches_1000_rmbounds.npy')
rmBounds = (rmBounds[0], rmBounds[1], rmBounds[2])

scCodes, icaCodes, v1Simple, v1cMean, angles = pr.responses(patches, scBasis, SPARSITY_PARAMETER, NON_NEGATIVE, icaFilters)
pcaTransformed, v1C = pr.responsesPCAV1C(patches)

patchesReconstV1 = pr.reconstructV1(v1Simple)
patchesReconstPCA = pr.reconstructPCA(pcaTransformed, v1cMean, angles)
patchesReconstSC = pr.reconstruct(scBasis, scCodes, v1cMean, angles)
patchesReconstICA = pr.reconstruct(icaFilters, icaCodes, v1cMean, angles)
'''
for i in range(patchesReconstV1.shape[0]):
	f, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, sharey = True)

	ax1.imshow(patchesOrig[i], cmap = 'gray', interpolation = 'none')
	ax1.title.set_text('Original patch')

	ax2.imshow(patches[i], cmap = 'gray', interpolation = 'none')
	ax2.title.set_text('Modified patch')

	ax3.imshow(patchesReconstV1[i], cmap = 'gray', interpolation = 'none')
	ax3.title.set_text('V1')

	ax4.imshow(patchesReconstSC[i], cmap = 'gray', interpolation = 'none')
	ax4.title.set_text('SC')

	plt.savefig('pred/' + str(i) + '.png')
	plt.close()
'''
